import os
import shutil
import random
import argparse
import concurrent.futures
import time
from tqdm import tqdm
import multiprocessing

# 设置线程池大小为CPU核心数的2倍
MAX_WORKERS = multiprocessing.cpu_count() * 2


def create_directory_structure(output_dir):
    """创建YOLO数据集所需的目录结构"""
    # 创建主目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 创建images和labels目录
    for folder in ['images', 'labels']:
        for split in ['train', 'val']:
            dir_path = os.path.join(output_dir, folder, split)
            os.makedirs(dir_path, exist_ok=True)
    
    print(f"已创建目录结构: {output_dir}\n- images/train\n- images/val\n- labels/train\n- labels/val")


def check_input_files(input_dir):
    """检查输入目录中的文件是否符合要求"""
    image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff'}
    label_extensions = {'.txt'}
    
    images = {}  # 使用字典直接存储文件名和路径
    labels = {}  # 使用字典直接存储文件名和路径
    
    # 使用os.scandir更快地遍历目录
    with os.scandir(input_dir) as entries:
        for entry in entries:
            if entry.is_file(follow_symlinks=False):
                name, ext = os.path.splitext(entry.name)
                if ext.lower() in image_extensions:
                    images[name] = entry.name
                elif ext.lower() in label_extensions:
                    labels[name] = entry.name
    
    # 检查图片和标签是否匹配
    image_names = set(images.keys())
    label_names = set(labels.keys())
    
    # 找出没有对应标签的图片
    images_without_labels = image_names - label_names
    if images_without_labels:
        print(f"警告: 发现 {len(images_without_labels)} 张图片没有对应的标签文件")
        if len(images_without_labels) < 10:
            print(f"未匹配的图片: {', '.join(images_without_labels)}")
    
    # 找出没有对应图片的标签
    labels_without_images = label_names - image_names
    if labels_without_images:
        print(f"警告: 发现 {len(labels_without_images)} 个标签文件没有对应的图片")
        if len(labels_without_images) < 10:
            print(f"未匹配的标签: {', '.join(labels_without_images)}")
    
    # 计算匹配的数量
    matched_count = len(image_names & label_names)
    print(f"找到 {matched_count} 对匹配的图片和标签")
    
    return images, labels, matched_count


def can_use_hardlink(src, dst):
    """检测是否可以使用硬链接

    检查条件:
    1. 源文件和目标目录是否在同一文件系统
    2. 系统是否支持硬链接
    """
    try:
        # 检查源文件和目标目录是否在同一文件系统
        src_dev = os.stat(src).st_dev
        dst_dev = os.stat(os.path.dirname(dst)).st_dev
        if src_dev != dst_dev:
            print(f"无法使用硬链接: 源文件和目标目录在不同文件系统")
            return False
        
        # 尝试创建一个临时硬链接来测试
        test_link = os.path.join(os.path.dirname(dst), f"temp_link_test_{os.getpid()}")
        os.link(src, test_link)
        os.unlink(test_link)
        return True
    except OSError as e:
        if e.errno == 1:  # 不支持操作
            print("无法使用硬链接: 系统不支持")
        elif e.errno == 13:  # 权限不足
            print("无法使用硬链接: 权限不足")
        else:
            print(f"无法使用硬链接: {str(e)}")
        return False
    except Exception as e:
        print(f"无法使用硬链接: {str(e)}")
        return False


def copy_file(src, dst, use_hardlink=False):
    """复制单个文件，可选使用硬链接"""
    try:
        if use_hardlink:
            os.link(src, dst)
            return True
        else:
            # 使用copyfile而不是copy2，因为我们通常不需要保留元数据
            shutil.copyfile(src, dst)
            return True
    except Exception as e:
        print(f"操作失败: {src} -> {dst}, 错误: {str(e)}")
        return False


def split_dataset(images, labels, train_ratio, input_dir, output_dir, use_hardlink=False):
    """将数据集分割为训练集和验证集"""
    # 提取匹配的文件名
    common_names = list(set(images.keys()) & set(labels.keys()))
    
    # 打乱顺序
    random.shuffle(common_names)
    
    # 计算分割点
    train_count = int(len(common_names) * train_ratio)
    val_count = len(common_names) - train_count
    
    print(f"数据集分割: {train_count} 张图片用于训练, {val_count} 张图片用于验证")
    
    # 准备训练集复制任务
    train_tasks = []
    for name in common_names[:train_count]:
        # 图片复制任务
        src_image = os.path.join(input_dir, images[name])
        dst_image = os.path.join(output_dir, 'images', 'train', images[name])
        train_tasks.append((src_image, dst_image))
        
        # 标签复制任务
        src_label = os.path.join(input_dir, labels[name])
        dst_label = os.path.join(output_dir, 'labels', 'train', labels[name])
        train_tasks.append((src_label, dst_label))
    
    # 准备验证集复制任务
    val_tasks = []
    for name in common_names[train_count:]:
        # 图片复制任务
        src_image = os.path.join(input_dir, images[name])
        dst_image = os.path.join(output_dir, 'images', 'val', images[name])
        val_tasks.append((src_image, dst_image))
        
        # 标签复制任务
        src_label = os.path.join(input_dir, labels[name])
        dst_label = os.path.join(output_dir, 'labels', 'val', labels[name])
        val_tasks.append((src_label, dst_label))
    
    # 复制训练集文件
    start_time = time.time()
    print(f"正在{'创建硬链接' if use_hardlink else '复制'}训练集文件 (使用 {min(MAX_WORKERS, len(train_tasks))} 个线程)...")
    
    with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        # 使用tqdm显示进度
        futures = {executor.submit(copy_file, src, dst, use_hardlink): (src, dst) for src, dst in train_tasks}
        for future in tqdm(concurrent.futures.as_completed(futures), total=len(futures), desc="训练集进度"):
            pass
    
    # 复制验证集文件
    print(f"正在{'创建硬链接' if use_hardlink else '复制'}验证集文件 (使用 {min(MAX_WORKERS, len(val_tasks))} 个线程)...")
    
    with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        # 使用tqdm显示进度
        futures = {executor.submit(copy_file, src, dst, use_hardlink): (src, dst) for src, dst in val_tasks}
        for future in tqdm(concurrent.futures.as_completed(futures), total=len(futures), desc="验证集进度"):
            pass
    
    end_time = time.time()
    print(f"文件{'硬链接' if use_hardlink else '复制'}完成，耗时: {end_time - start_time:.2f} 秒")


def create_data_yaml(output_dir, class_names=None):
    """创建YOLO训练所需的data.yaml文件"""
    if class_names is None:
        class_names = ['class0', 'class1', 'class2']  # 默认类别
        print("警告: 未提供类别名称，使用默认类别 ['class0', 'class1', 'class2']")
        print("提示: 请根据您的实际数据集修改data.yaml文件中的类别名称")
    
    yaml_content = f"""path: {output_dir}\n\n# 训练集和验证集的相对路径\ntrain: images/train\nval: images/val\n\n# 类别数量\nnc: {len(class_names)}\n\n# 类别名称\nnames: {class_names}\n"""
    
    yaml_path = os.path.join(output_dir, 'data.yaml')
    with open(yaml_path, 'w', encoding='utf-8') as f:
        f.write(yaml_content)
    
    print(f"已创建data.yaml文件: {yaml_path}")


def get_user_input(prompt, default=None, is_directory=True):
    """获取用户输入，并进行验证"""
    while True:
        user_input = input(prompt)
        if not user_input and default is not None:
            return default
        
        if is_directory:
            if os.path.exists(user_input):
                return user_input
            else:
                print(f"错误: 目录不存在: {user_input}")
        else:
            return user_input


def main():
    # 全局变量设置
    global MAX_WORKERS
    
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='准备YOLOv11训练数据集')
    parser.add_argument('--input', '-i', help='输入文件夹路径，包含图片和标签文件')
    parser.add_argument('--output', '-o', help='输出文件夹路径')
    parser.add_argument('--train-ratio', '-r', type=float, default=0.8, help='训练集比例 (默认: 0.8)')
    parser.add_argument('--classes', '-c', nargs='+', help='类别名称列表，例如: --classes person car bicycle')
    parser.add_argument('--workers', '-w', type=int, help=f'线程数量 (默认: CPU核心数×2 = {MAX_WORKERS})')
    parser.add_argument('--use-hardlink', '-l', action='store_true', help='尝试使用硬链接代替文件复制 (可以提高速度并节省空间)')
    parser.add_argument('--no-hardlink', action='store_true', help='禁用硬链接，强制使用文件复制')
    args = parser.parse_args()

    if args.workers:
        MAX_WORKERS = max(1, args.workers)
        print(f"使用自定义线程数量: {MAX_WORKERS}")
    else:
        print(f"使用默认线程数量: {MAX_WORKERS} (CPU核心数×2)")

    # 如果没有通过命令行参数提供，则交互式获取输入
    if not args.input:
        print("欢迎使用YOLOv11数据集准备工具!")
        input_dir = get_user_input("请输入包含图片和标签的文件夹路径: ")
    else:
        input_dir = args.input
    
    if not args.output:
        output_dir = get_user_input("请输入输出文件夹路径: ")
    else:
        output_dir = args.output
    
    train_ratio = args.train_ratio
    
    if not args.classes:
        classes_input = input("请输入类别名称 (用空格分隔，默认: class0 class1 class2): ")
        class_names = classes_input.split() if classes_input else None
    else:
        class_names = args.classes

    # 确定是否使用硬链接
    use_hardlink = False
    if args.no_hardlink:
        print("已禁用硬链接，将使用文件复制")
        use_hardlink = False
    elif args.use_hardlink:
        print("尝试使用硬链接...")
        # 获取一个示例文件来测试硬链接
        sample_files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
        if sample_files:
            sample_file = os.path.join(input_dir, sample_files[0])
            sample_dst = os.path.join(output_dir, 'test_hardlink')
            use_hardlink = can_use_hardlink(sample_file, sample_dst)
        else:
            print("输入目录中没有文件，无法测试硬链接")
            use_hardlink = False
    else:
        # 自动模式：询问用户是否要尝试使用硬链接
        if os.name != 'nt':  # 非Windows系统默认支持硬链接
            response = input("是否尝试使用硬链接提高速度并节省空间? (y/n, 默认: y): ").lower()
            if response != 'n':
                print("尝试使用硬链接...")
                # 获取一个示例文件来测试硬链接
                sample_files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
                if sample_files:
                    sample_file = os.path.join(input_dir, sample_files[0])
                    sample_dst = os.path.join(output_dir, 'test_hardlink')
                    use_hardlink = can_use_hardlink(sample_file, sample_dst)
                else:
                    print("输入目录中没有文件，无法测试硬链接")
                    use_hardlink = False
        else:
            print("Windows系统上硬链接支持有限，默认使用文件复制")
            use_hardlink = False

    # 检查训练集比例是否有效
    if not (0 < train_ratio < 1):
        print(f"错误: 训练集比例必须在0和1之间: {train_ratio}")
        return
    
    # 创建输出目录结构
    create_directory_structure(output_dir)
    
    # 检查输入文件
    start_time = time.time()
    images, labels, matched_count = check_input_files(input_dir)
    check_time = time.time() - start_time
    print(f"文件检查完成，耗时: {check_time:.2f} 秒")
    
    if matched_count == 0:
        print("错误: 没有找到匹配的图片和标签文件")
        return
    
    # 分割数据集
    split_dataset(images, labels, train_ratio, input_dir, output_dir, use_hardlink)
    
    # 创建data.yaml文件
    create_data_yaml(output_dir, class_names)
    
    print(f"\n数据集准备完成！\n输出目录: {output_dir}")
    print("请检查data.yaml文件，确保类别名称和数量正确。")


if __name__ == '__main__':
    main()